# -*- coding: utf-8 -*-
# file: data_utils.py
# author: songyouwei <youwei0314@gmail.com>
# Copyright (C) 2018. All Rights Reserved.


'''
    tokenizer:
    1.berttokenizer自建词库，能够直接将输入的字符串转换为index
    2.自己写的tokenizer首先读入所有的字符，建立word2index字典，然后再实现截取和padding
'''
import os
import pickle  #对象序列化
import numpy as np
import torch
from torch.utils.data import Dataset  #数据打包的类，结果传入到DataLoader中
from transformers import BertTokenizer
import utils.data_utils_noBert

#padding或者截取，sequence是index序列,所有词嵌入都需要用到
def pad_and_truncate(sequence, maxlen, dtype='int64', padding='post', truncating='post', value=0):
    x = (np.ones(maxlen) * value).astype(dtype)  #建立最大长度的字符数组，多的补1
    if truncating == 'pre':  #取截取的部分
        trunc = sequence[-maxlen:]  #从后往前截取
    else:
        trunc = sequence[:maxlen]  #从前往后截取
    trunc = np.asarray(trunc, dtype=dtype)  #转换为字符串数组

    if padding == 'post':
        x[:len(trunc)] = trunc #将从后往前截取的部分作为返回值
    else:
        x[-len(trunc):] = trunc #将从前往后截取x的部分作为返回值
    return x


class Tokenizer4Bert:  #Bert使用的分词器，给出预训练模型的名字和最大长度，加入cls等
    '''
           BERT_PATH = 'D:/thesis/code/dataset/bert/bert-base-uncased'
           self.tokenizer = BertTokenizer.from_pretrained(BERT_PATH)  #加载本地的分词器
    '''
    def __init__(self, max_seq_len, pretrained_bert_name):
        self.tokenizer = BertTokenizer.from_pretrained(pretrained_bert_name)  #获取Bert分词器
        self.max_seq_len = max_seq_len

    def text_to_sequence(self, text, reverse=False, padding='post', truncating='post'): #将文本转换为单词序列，包装了一层，里面使用的是berttokenizer
        sequence = self.tokenizer.convert_tokens_to_ids(self.tokenizer.tokenize(text)) #将bert解析出的word序列转换为tokenizer中的id
        if len(sequence) == 0:
            sequence = [0]
        if reverse:  #反转词语序列
            sequence = sequence[::-1]
        return pad_and_truncate(sequence, self.max_seq_len, padding=padding, truncating=truncating)  #返回padding和截取之后的结果


class ABSADataset(Dataset):  #重写了dataset父类，创建自己的数据集
    def __init__(self, fname, tokenizer):  #读入原始数据,fname是配置文件中给出的数据集名字(执行的时候给的是opt.dataset_file['train'])
        fin = open(fname, 'r', encoding='utf-8', newline='\n', errors='ignore')  #打开文件
        lines = fin.readlines()   #读取文件，并且获取文本内容的行数,lattop有6984行
        fin.close()

        fin = open(fname+'.graph', 'rb')  #以二进制只读形式打开.graph文件（句法依赖？）
        idx2graph = pickle.load(fin)  #idx2graph:dict,2328
        fin.close()

        all_data = []
        for i in range(0, len(lines), 3):  #步长为3，因为读取的数据第一行是句子，第二行是方面，第三行是这个方面的极性
            text_left, _, text_right = [s.lower().strip() for s in lines[i].partition("$T$")]  #得到方面词的左，右侧文本和方面词,partition() 方法用来根据指定的分隔符将字符串进行分割。S为三部分字符串
            aspect = lines[i + 1].lower().strip()  #得到方面词,strip() 方法用于移除字符串头尾指定的字符（默认为空格或换行符）或字符序列。
            polarity = lines[i + 2].strip()  #得到关于这个方面的极性,strip() 方法用于移除字符串头尾指定的字符（默认为空格或换行符）或字符序列。

            #处理一个句子得到的属性，如果是Bert模型的话，就采用tokenizer4bert构造分词器，否则用tokenizer
            text_indices = tokenizer.text_to_sequence(text_left + " " + aspect + " " + text_right)  #将原文本得到index序列
            context_indices = tokenizer.text_to_sequence(text_left + " " + text_right) #将上下文得到词典index序列

            left_indices = tokenizer.text_to_sequence(text_left)  #左侧文本词典index序列
            left_with_aspect_indices = tokenizer.text_to_sequence(text_left + " " + aspect)  #带上方面词的左侧文本词典index序列
            right_indices = tokenizer.text_to_sequence(text_right, reverse=True)  ##右侧文本词典index序列，翻转，
            right_with_aspect_indices = tokenizer.text_to_sequence(aspect + " " + text_right, reverse=True)  #带上方面词的右侧侧文本词典index序列

            aspect_indices = tokenizer.text_to_sequence(aspect)  #方面词转换为词典index
            left_len = np.sum(left_indices != 0)  #计算左侧文本长度
            aspect_len = np.sum(aspect_indices != 0)  #方面短语的长度，即方面短语由几个词组成
            aspect_boundary = np.asarray([left_len, left_len + aspect_len - 1], dtype=np.int64) #方面词的上下界限
            polarity = int(polarity) + 1 #0是消极，1是中性，2是积极

            text_len = np.sum(text_indices != 0)  #对np中非零元素个数进行求和，也就是单词长度
            #将bert类型的方面词连接到bert类型的正文之后
            concat_bert_indices = tokenizer.text_to_sequence('[CLS] ' + text_left + " " + aspect + " " + text_right + ' [SEP] ' + aspect + " [SEP]")
            concat_segments_indices = [0] * (text_len + 2) + [1] * (aspect_len + 1) #text_left + " " + aspect + " " + text_right为0，“ ”+aspect为1
            concat_segments_indices = pad_and_truncate(concat_segments_indices, tokenizer.max_seq_len)
            #正文的bert类型index
            text_bert_indices = tokenizer.text_to_sequence("[CLS] " + text_left + " " + aspect + " " + text_right + " [SEP]")
            aspect_bert_indices = tokenizer.text_to_sequence("[CLS] " + aspect + " [SEP]")
            #i=0代表句子，i=1代表方面词，i=2代表极性，用constant方式填充，dependency_graph=(85,85)
            dependency_graph = np.pad(idx2graph[i], \
                ((0,tokenizer.max_seq_len-idx2graph[i].shape[0]),(0,tokenizer.max_seq_len-idx2graph[i].shape[0])), 'constant')

            data = {   #字典,index就是词语在字典中的索引
                'concat_bert_indices': concat_bert_indices,
                'concat_segments_indices': concat_segments_indices,
                'text_bert_indices': text_bert_indices,
                'aspect_bert_indices': aspect_bert_indices,
                'text_indices': text_indices,
                'context_indices': context_indices,
                'left_indices': left_indices,
                'left_with_aspect_indices': left_with_aspect_indices,
                'right_indices': right_indices,
                'right_with_aspect_indices': right_with_aspect_indices,
                'aspect_indices': aspect_indices,  #得到方面词转换为index序列
                'aspect_boundary': aspect_boundary,  #
                'dependency_graph': dependency_graph,  #依赖图
                'polarity': polarity, #极性
            }

            all_data.append(data)  #所有数据，每个item都是上面的字典
        self.data = all_data

    def __getitem__(self, index):
        return self.data[index]  #重写[]运算符

    def __len__(self):
        return len(self.data)  #返回data的长度
